# data_sources/baostock_client.py
from datetime import datetime

import baostock as bs
import pandas as pd
import time
import logging
from typing import List, Dict, Optional

from stock_data_system.config.database import BaostockConfig

class BaostockClient:
    """Baostock数据客户端"""

    def __init__(self, config: BaostockConfig):
        self.config = config
        self.lg = None
        self.logger = logging.getLogger(__name__)
        self._login()

    def _login(self):
        """登录Baostock"""
        try:
            self.lg = bs.login()
            if self.lg.error_code == '0':
                self.logger.info("Baostock登录成功")
            else:
                self.logger.error(f"Baostock登录失败: {self.lg.error_msg}")
        except Exception as e:
            self.logger.error(f"Baostock登录异常: {e}")

    def __del__(self):
        """退出时自动登出"""
        if self.lg:
            bs.logout()
            self.logger.info("Baostock已登出")

    def get_stock_basic(self) -> pd.DataFrame:
        """获取股票基本信息"""
        try:
            rs = bs.query_stock_basic()
            if rs.error_code != '0':
                self.logger.error(f"获取股票列表失败: {rs.error_msg}")
                return pd.DataFrame()

            data_list = []
            while (rs.error_code == '0') and rs.next():
                data_list.append(rs.get_row_data())

            if not data_list:
                self.logger.warning("未获取到股票数据")
                return pd.DataFrame()

            df = pd.DataFrame(data_list, columns=rs.fields)

            # 标准化列名
            column_mapping = {
                'code': 'symbol',
                'code_name': 'name',
                'ipoDate': 'listing_date',
                'outDate': 'delisting_date',
                'type': 'type',
                'status': 'status'
            }
            df = df.rename(columns=column_mapping)

            # 添加市场信息
            df['market'] = df['symbol'].apply(self._get_market_from_symbol)
            df['country'] = 'CN'
            df['currency'] = 'CNY'

            # 过滤掉已退市股票
            df = df[df['status'] == '1']

            self.logger.info(f"获取到 {len(df)} 只股票基本信息")
            return df

        except Exception as e:
            self.logger.error(f"获取股票基本信息异常: {e}")
            return pd.DataFrame()

    def get_daily_data(self, symbol: str, start_date: str, end_date: str) -> pd.DataFrame:
        """获取日线数据"""
        try:
            # 确保股票代码格式正确
            formatted_symbol = self._format_symbol(symbol)
            if not formatted_symbol:
                self.logger.error(f"股票代码格式不正确: {symbol}")
                return pd.DataFrame()

            # 查询历史K线数据
            rs = bs.query_history_k_data_plus(
                formatted_symbol,
                "date,code,open,high,low,close,volume,amount,turn,pctChg,peTTM,pbMRQ",
                start_date=start_date,
                end_date=end_date,
                frequency="d",
                adjustflag="3"  # 后复权
            )

            if rs.error_code != '0':
                self.logger.error(f"获取日线数据失败 {formatted_symbol}: {rs.error_msg}")
                return pd.DataFrame()

            data_list = []
            while (rs.error_code == '0') and rs.next():
                data_list.append(rs.get_row_data())

            if not data_list:
                self.logger.info(f"{formatted_symbol} 没有找到数据")
                return pd.DataFrame()

            df = pd.DataFrame(data_list, columns=rs.fields)

            # 数据类型转换
            numeric_columns = ['open', 'high', 'low', 'close', 'volume', 'amount', 'turn', 'pctChg', 'peTTM', 'pbMRQ']
            for col in numeric_columns:
                if col in df.columns:
                    df[col] = pd.to_numeric(df[col], errors='coerce')

            # 标准化列名和数据格式
            df = self._normalize_daily_data(df)

            # 添加简化股票代码
            df['symbol'] = self._extract_stock_code(formatted_symbol)

            # 控制请求频率
            time.sleep(self.config.request_delay)

            self.logger.info(f"获取 {formatted_symbol} 日线数据 {len(df)} 条")
            return df

        except Exception as e:
            self.logger.error(f"获取日线数据异常 {symbol}: {e}")
            return pd.DataFrame()

    def _format_symbol(self, symbol: str) -> str:
        """格式化股票代码为Baostock要求的格式"""
        if not symbol:
            return ""

        # 如果已经是正确格式
        if '.' in symbol and len(symbol) >= 9:
            return symbol

        # 纯数字代码添加市场前缀
        if symbol.isdigit():
            if symbol.startswith('6'):
                return f"sh.{symbol}"
            elif symbol.startswith('0') or symbol.startswith('3'):
                return f"sz.{symbol}"
            elif symbol.startswith('4') or symbol.startswith('8'):
                return f"bj.{symbol}"

        # 其他格式转换
        symbol_upper = symbol.upper()
        if symbol_upper.startswith(('SH', 'SZ', 'BJ')):
            market = symbol_upper[:2].lower()
            code = symbol_upper[2:]
            return f"{market}.{code}"

        self.logger.warning(f"无法识别的股票代码格式: {symbol}")
        return ""

    def _get_market_from_symbol(self, symbol: str) -> str:
        """从股票代码提取市场信息"""
        if not symbol or '.' not in symbol:
            return 'UNKNOWN'

        market_prefix = symbol.split('.')[0]
        market_map = {'sh': 'SH', 'sz': 'SZ', 'bj': 'BJ'}
        return market_map.get(market_prefix, 'UNKNOWN')

    def _extract_stock_code(self, symbol: str) -> str:
        """提取纯股票代码（去掉市场前缀）"""
        if not symbol or '.' not in symbol:
            return symbol
        return symbol.split('.')[1]

    def _normalize_daily_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """标准化日线数据格式"""
        if df.empty:
            return df

        column_mapping = {
            'date': 'trade_date',
            'code': 'symbol',
            'open': 'open_price',
            'high': 'high_price',
            'low': 'low_price',
            'close': 'close_price',
            'volume': 'volume',
            'amount': 'turnover',
            'pctChg': 'change_rate',
            'turn': 'turnover_rate',
            'peTTM': 'pe_ratio',
            'pbMRQ': 'pb_ratio'
        }

        # 只重命名存在的列
        df = df.rename(columns={col: column_mapping[col] for col in df.columns if col in column_mapping})

        # 确保日期格式正确
        if 'trade_date' in df.columns:
            df['trade_date'] = pd.to_datetime(df['trade_date']).dt.strftime('%Y-%m-%d')

        return df

    def test_connection(self) -> bool:
        """测试Baostock连接"""
        try:
            if not self.lg:
                self._login()

            # 测试获取上证指数信息
            rs = bs.query_stock_basic("sh.000001")
            return rs.error_code == '0'

        except Exception as e:
            self.logger.error(f"Baostock连接测试失败: {e}")
            return False

    def get_financial_data(self, symbol: str, year: int = None, quarter: int = None) -> pd.DataFrame:
        """获取财务数据"""
        try:
            formatted_symbol = self._format_symbol(symbol)
            if not formatted_symbol:
                self.logger.error(f"股票代码格式不正确: {symbol}")
                return pd.DataFrame()

            # 如果没有指定年份，使用最近一年
            if year is None:
                year = datetime.now().year - 1

            # 获取盈利能力数据
            profit_data = self._get_profit_data(formatted_symbol, year, quarter)
            # 获取成长能力数据
            growth_data = self._get_growth_data(formatted_symbol, year, quarter)
            # 获取营运能力数据
            operation_data = self._get_operation_data(formatted_symbol, year, quarter)
            # 获取偿债能力数据
            balance_data = self._get_balance_data(formatted_symbol, year, quarter)

            # 合并财务数据
            financial_data = {**profit_data, **growth_data, **operation_data, **balance_data}

            if not financial_data:
                self.logger.info(f"{symbol} 没有找到财务数据")
                return pd.DataFrame()

            # 添加基本信息
            financial_data['symbol'] = symbol
            financial_data['report_date'] = self._get_report_date(year, quarter)
            financial_data['report_type'] = self._get_report_type(quarter)
            financial_data['data_source'] = 'baostock'

            df = pd.DataFrame([financial_data])

            # 控制请求频率
            time.sleep(self.config.request_delay)

            self.logger.info(f"获取 {symbol} 财务数据成功")
            return df

        except Exception as e:
            self.logger.error(f"获取财务数据异常 {symbol}: {e}")
            return pd.DataFrame()

    def _get_profit_data(self, symbol: str, year: int, quarter: int) -> Dict:
        """获取盈利能力数据"""
        try:
            rs = bs.query_profit_data(code=symbol, year=year, quarter=quarter)
            if rs.error_code != '0' or not rs.next():
                return {}

            data = rs.get_row_data()
            return {
                'roe': self._safe_float(data[1]),  # 净资产收益率
                'net_profit': self._safe_float(data[2]),  # 净利润
                'eps': self._safe_float(data[3]),  # 每股收益
                'revenue': self._safe_float(data[4]),  # 营业收入
                'total_profit': self._safe_float(data[5]),  # 利润总额
                'operating_profit': self._safe_float(data[6]),  # 营业利润
                'gross_margin': self._safe_float(data[13]),  # 毛利率
                'net_margin': self._safe_float(data[14]),  # 净利率
            }
        except Exception as e:
            self.logger.warning(f"获取盈利能力数据失败 {symbol}: {e}")
            return {}

    def _get_growth_data(self, symbol: str, year: int, quarter: int) -> Dict:
        """获取成长能力数据"""
        try:
            rs = bs.query_growth_data(code=symbol, year=year, quarter=quarter)
            if rs.error_code != '0' or not rs.next():
                return {}

            data = rs.get_row_data()
            return {
                'revenue_growth_rate': self._safe_float(data[1]),  # 营业收入增长率
                'profit_growth_rate': self._safe_float(data[2]),  # 利润总额增长率
                'asset_growth_rate': self._safe_float(data[3]),  # 总资产增长率
                'equity_growth_rate': self._safe_float(data[4]),  # 净资产增长率
                'eps_growth_rate': self._safe_float(data[5]),  # 每股收益增长率
            }
        except Exception as e:
            self.logger.warning(f"获取成长能力数据失败 {symbol}: {e}")
            return {}

    def _get_operation_data(self, symbol: str, year: int, quarter: int) -> Dict:
        """获取营运能力数据"""
        try:
            rs = bs.query_operation_data(code=symbol, year=year, quarter=quarter)
            if rs.error_code != '0' or not rs.next():
                return {}

            data = rs.get_row_data()
            return {
                'asset_turnover': self._safe_float(data[1]),  # 总资产周转率
                'inventory_turnover': self._safe_float(data[4]),  # 存货周转率
                'receivable_turnover': self._safe_float(data[7]),  # 应收账款周转率
            }
        except Exception as e:
            self.logger.warning(f"获取营运能力数据失败 {symbol}: {e}")
            return {}

    def _get_balance_data(self, symbol: str, year: int, quarter: int) -> Dict:
        """获取偿债能力数据"""
        try:
            rs = bs.query_balance_data(code=symbol, year=year, quarter=quarter)
            if rs.error_code != '0' or not rs.next():
                return {}

            data = rs.get_row_data()
            return {
                'total_assets': self._safe_float(data[1]),  # 总资产
                'total_liabilities': self._safe_float(data[2]),  # 总负债
                'equity': self._safe_float(data[3]),  # 净资产
                'current_assets': self._safe_float(data[6]),  # 流动资产
                'current_liabilities': self._safe_float(data[8]),  # 流动负债
                'asset_liability_ratio': self._safe_float(data[13]),  # 资产负债率
            }
        except Exception as e:
            self.logger.warning(f"获取偿债能力数据失败 {symbol}: {e}")
            return {}

    def get_dividend_data(self, symbol: str, year: int = None) -> pd.DataFrame:
        """获取分红数据"""
        try:
            formatted_symbol = self._format_symbol(symbol)
            if not formatted_symbol:
                self.logger.error(f"股票代码格式不正确: {symbol}")
                return pd.DataFrame()

            # 如果没有指定年份，获取最近5年的数据
            if year is None:
                current_year = datetime.now().year
                years = list(range(current_year - 4, current_year + 1))
            else:
                years = [year]

            all_dividends = []

            for y in years:
                rs = bs.query_dividend_data(code=formatted_symbol, year=y, yearType="report")

                if rs.error_code == '0':
                    while rs.next():
                        data = rs.get_row_data()
                        dividend_info = {
                            'symbol': symbol,
                            'dividend_date': self._parse_date(data[3]),  # 分红日期
                            'ex_dividend_date': self._parse_date(data[4]),  # 除权除息日
                            'dividend_per_share': self._safe_float(data[5]),  # 每股股利
                            'dividend_ratio': self._safe_float(data[6]),  # 股息率
                            'update_date': datetime.now().strftime('%Y-%m-%d'),
                            'data_source': 'baostock'
                        }
                        all_dividends.append(dividend_info)

                time.sleep(self.config.request_delay)  # 控制请求频率

            if not all_dividends:
                self.logger.info(f"{symbol} 没有找到分红数据")
                return pd.DataFrame()

            df = pd.DataFrame(all_dividends)
            self.logger.info(f"获取 {symbol} 分红数据 {len(df)} 条")
            return df

        except Exception as e:
            self.logger.error(f"获取分红数据异常 {symbol}: {e}")
            return pd.DataFrame()

    def _safe_float(self, value):
        """安全转换为float"""
        try:
            return float(value) if value and value != '' else None
        except (ValueError, TypeError):
            return None

    def _parse_date(self, date_str):
        """解析日期字符串"""
        if not date_str or date_str == '':
            return None
        try:
            return datetime.strptime(date_str, '%Y-%m-%d').strftime('%Y-%m-%d')
        except ValueError:
            return None

    def _get_report_date(self, year: int, quarter: int) -> str:
        """生成报告日期"""
        if quarter == 1:
            return f"{year}-03-31"
        elif quarter == 2:
            return f"{year}-06-30"
        elif quarter == 3:
            return f"{year}-09-30"
        elif quarter == 4:
            return f"{year}-12-31"
        else:
            return f"{year}-12-31"

    def _get_report_type(self, quarter: int) -> str:
        """获取报告类型"""
        if quarter == 1:
            return "Q1"
        elif quarter == 2:
            return "HY"  # 半年度
        elif quarter == 3:
            return "Q3"
        elif quarter == 4:
            return "FY"  # 年度
        else:
            return "FY"